# -*- coding: utf-8 -*-
"""
Created on Mon Oct  8 14:32:52 2018

@author: luolei

交叉映射相关性计算
"""
import pandas as pd
import numpy as np


def cross_correlation(series_a, series_b, d, eps = 1e-6):
	"""
	在延迟为d上的互相关分析, series_a, 移动series_b: d > 0向右移，d < 0向左移
	:param series_a: np.ndarray, 目标变量, shape = (-1,)
	:param series_b: np.ndarray, 外生变量, shape = (-1,)
	:param d: time delay, 延迟阶数
	:return:
	"""
	try:
		assert type(series_a) == type(series_b) == np.ndarray
		assert len(series_a.shape) == len(series_b.shape) == 1
		assert len(series_a) == len(series_b)
	except Exception:
		raise ValueError('输入数据格式不正确')
	
	series_x = series_a.copy()
	series_y = series_b.copy()
	
	len_y = len(series_y)
	series_y = np.hstack((series_y[-(d % len_y):], series_y[: -(d % len_y)]))
	
	mean_x, mean_y = np.mean(series_x), np.mean(series_y)
	numerator = np.sum((series_x - mean_x) * (series_y - mean_y))
	denominator = np.sqrt(np.sum(np.power((series_x - mean_x), 2))) * np.sqrt(
		np.sum(np.power((series_y - mean_y), 2)))
	
	return numerator / (denominator + eps)


def mean_ccf(series_a, series_b, d, seg_len, start_locs):
	"""平均CCF值"""
	series_list = []
	try:
		assert start_locs[0] + d >= 0
	except Exception:
		raise ValueError('start locs起始位置不能被分析')
	
	try:
		assert start_locs[-1] + d <= len(series_a) - 1
	except Exception:
		raise ValueError('start locs末尾位置不能被分析')
	
	for start_loc in start_locs:
		series_list.append(
			[
				series_a.copy()[start_loc: start_loc + seg_len],
				series_b.copy()[(start_loc + d): (start_loc + seg_len + d)]])
	
	series_list = pd.DataFrame(series_list, columns = ['series_a', 'series_b'])
	series_list['ccf_value'] = series_list.apply(lambda x: cross_correlation(x[0], x[1], 0), axis = 1)
	ccf = np.mean(list(series_list['ccf_value']))
	
	return ccf


def time_delay_ccf_test(series_a, series_b, d_list, seg_len, start_locs):
	"""
	含时滞计算的ccf
	:param series_a: array
	:param series_b: array
	:param d_list: list or array
	:param seg_len: int
	:param start_locs: list or array
	:return: time_delay_ccf, dict
	"""
	time_delay_ccf = {}
	for d in d_list:
		time_delay_ccf[int(d)] = mean_ccf(series_a, series_b, d, seg_len, start_locs)
	
	return time_delay_ccf


